import cv2
import numpy as np
import matplotlib.pyplot as plt
import os

dir = '../../../../../large_data/pic/watershed/'
name = 'poker_cards.png'
# dir = '../../../../../large_data/pic/'
# name = 'dog_bird.png'
path = os.path.join(dir, name)
img = cv2.imread(path, cv2.IMREAD_COLOR)

def cv_show(name , img):
    cv2.imshow(name , img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

#预处理
gray = cv2.cvtColor(img , cv2.COLOR_BGR2GRAY)
ret , img_thr = cv2.threshold(gray , 0 , 255 , cv2.THRESH_BINARY + cv2.THRESH_OTSU)
ksize = np.ones((3 , 3) , np.uint8)
img_mor = cv2.morphologyEx(img_thr , cv2.MORPH_OPEN , ksize , iterations=2)
# cv_show('img' , img_mor)

#膨胀找背景
img_dil = cv2.dilate(img_mor , ksize , iterations=2)
# cv_show('img' , img_dil)

#找出中心点

img_dis = cv2.distanceTransform(img_mor , 1 , 5)
img_dis = cv2.erode(img_dis , ksize , iterations=7)

plt.imshow(img_dis , 'gray')
plt.show()

#二值化
ret , img_thr2 = cv2.threshold(img_dis , 0.65*img_dis.max() , 255 , 0)

#找出未知区域
img_thr2 = np.uint8(img_thr2)
img_sub = cv2.subtract(img_mor , img_thr2)

# cv_show('img' , img_sub)

#贴标签
ret , maek1 = cv2.connectedComponents(img_thr2)
print(maek1.max())
maek = maek1+1
maek[img_sub == 255] = 0

#分水岭算法
img_copy = img.copy()
maek3 = cv2.watershed(img_copy , maek)
img_copy[maek3 == -1] = [0 , 0 ,255]

# cv_show('img' , img_copy)

plt.subplot(241) , plt.imshow(img , 'gray') , plt.title('img') , plt.xticks([]) , plt.yticks([])
plt.subplot(242) , plt.imshow(img_thr , 'gray') , plt.title('img_thr') , plt.xticks([]) , plt.yticks([])
plt.subplot(243) , plt.imshow(img_mor , 'gray') , plt.title('img_mor') , plt.xticks([]) , plt.yticks([])
plt.subplot(244) , plt.imshow(img_dil , 'gray') , plt.title('img_dil') , plt.xticks([]) , plt.yticks([])
plt.subplot(245) , plt.imshow(img_dis , 'gray') , plt.title('img_dis') , plt.xticks([]) , plt.yticks([])
plt.subplot(246) , plt.imshow(img_thr2, 'gray') , plt.title('img_thr2') , plt.xticks([]) , plt.yticks([])
plt.subplot(247) , plt.imshow(maek  , 'jet') , plt.title('maek ') , plt.xticks([]) , plt.yticks([])
plt.subplot(248) , plt.imshow(img_copy , 'gray') , plt.title('img_copy') , plt.xticks([]) , plt.yticks([])
plt.show()







